
#include "Kernel.h"
#include "Common/InteractionLogic.h"

//Allocates memory on gpu and copies required data.
inline float *copyToGpu(float *data, size_t count)
{
	size_t size = sizeof(float) * count;
	float *res;
	SAFE_CALL(cudaMalloc((void **)&res, size));
	SAFE_CALL(cudaMemcpy(res, data, size, cudaMemcpyHostToDevice));
	return res;
}

//Copies data from GPU and releases allocated GPU-memory.
inline void copyFromGpu(float *data_host, float *&data_dev, size_t count)
{
	size_t size = sizeof(float) * count;
	SAFE_CALL(cudaMemcpy(data_host, data_dev, size, cudaMemcpyDeviceToHost));
	SAFE_CALL(cudaFree(data_dev));
	data_dev = NULL;
}

__global__ void interactionKernel(float *px, float *py, float *pz,
								   float *ax, float *ay, float *az,
								   float *masses, size_t count)
{
	int n = blockIdx.x * blockDim.x + threadIdx.x;
	if (n < count)
	{
		//Copying current body state.
		float _px = px[n];
		float _py = py[n];
		float _pz = pz[n];
		float _mass = masses[n];
		float _ax = ax[n];
		float _ay = ay[n];
		float _az = az[n];
		//Interacting.
		for (int i = 0; i < count; i++)
			if (i != n)
				InteractionLogic::interactWith(_px, _py, _pz,
											   _ax, _ay, _az, _mass,
											   px[i], py[i], pz[i],
											   masses[i]);
		//Storing acceleration.
		ax[n] = _ax;
		ay[n] = _ay;
		az[n] = _az;
	}
	__syncthreads();
}

__global__ void updateKernel(float *px, float *py, float *pz,
							 float *vx, float *vy, float *vz,
							 float *ax, float *ay, float *az,
							 size_t count, float dt)
{
	int n = blockIdx.x * blockDim.x + threadIdx.x;
	if (n < count)
	{
		//Copying body's parameters.
		float _px = px[n];
		float _py = py[n];
		float _pz = pz[n];
		float _vx = vx[n];
		float _vy = vy[n];
		float _vz = vz[n];
		float _ax = ax[n];
		float _ay = ay[n];
		float _az = az[n];
		//Updating.
		InteractionLogic::update(_px, _py, _pz,
								 _vx, _vy, _vz,
								 _ax, _ay, _az,
								 dt);
		//Storing results.
		px[n] = _px;
		py[n] = _py;
		pz[n] = _pz;
		vx[n] = _vx;
		vy[n] = _vy;
		vz[n] = _vz;
	}
	__syncthreads();
}

void interactAndUpdate(float *px, float *py, float *pz,
					   float *vx, float *vy, float *vz,
					   float *ax, float *ay, float *az,
					   float *masses, size_t count, size_t block_size)
{
	//Copying positions
	float *px_dev = copyToGpu(px, count);
	float *py_dev = copyToGpu(py, count);
	float *pz_dev = copyToGpu(pz, count);
	//Copying velocities
	float *vx_dev = copyToGpu(vx, count);
	float *vy_dev = copyToGpu(vy, count);
	float *vz_dev = copyToGpu(vz, count);
	//Copying accelerations
	float *ax_dev = copyToGpu(ax, count);
	float *ay_dev = copyToGpu(ay, count);
	float *az_dev = copyToGpu(az, count);
	//Copying masses
	float *masses_dev = copyToGpu(masses, count);
	//Launching kernels.
	dim3 threads(block_size);
	dim3 grid(count / block_size + 1);
	interactionKernel<<<threads, grid>>>(px_dev, py_dev, pz_dev,
										 ax_dev, ay_dev, az_dev,
										 masses_dev, count);
	updateKernel<<<threads, grid>>>(px_dev, py_dev, pz_dev,
								    vx_dev, vy_dev, vz_dev,
								    ax_dev, ay_dev, az_dev,
								    count, DTIME);
	//Copying new possitions.
	copyFromGpu(px, px_dev, count);
	copyFromGpu(py, py_dev, count);
	copyFromGpu(pz, pz_dev, count);
	//Copying new velocities.
	copyFromGpu(vx, vx_dev, count);
	copyFromGpu(vy, vy_dev, count);
	copyFromGpu(vz, vz_dev, count);
	//Releasing accelerations and masses.
	SAFE_CALL(cudaFree(ax_dev));
	SAFE_CALL(cudaFree(ay_dev));
	SAFE_CALL(cudaFree(az_dev));
	SAFE_CALL(cudaFree(masses_dev));
}
